#ifndef AGGMG_HPP
#define AGGMG_HPP

#include "precond.hpp"

/// @brief 用于管理在多重网格中不同层间发生进程数收缩时的数据结构
class SHADOW
{
public:
    // 所有的主进程组成的通信子
    MPI_Comm masterComm = MPI_COMM_NULL;
    // 一个主进程和对应若干从属进程组成的通信子
    MPI_Comm subComm = MPI_COMM_NULL;
    // 主进程所需的数据
    // 接收数目，等同于本主进程对应多少个从属进程
    int num_recvs = 0;
    // 从属进程的进程号，注意这个进程号是原来的整个通信域的
    int* recv_pids = nullptr;
    // 从从属进程收集回来的消息应该放置的位置偏移
    int* recv_offsets = nullptr;

    // 从属进程所需的数据
    // 向谁发？即主进程的进程号，注意这个进程号是原来的整个通信域的
    int send_pid = MPI_PROC_NULL;
    // 要向主进程发送的消息长度
    int send_len = 0;

    /// @brief 析构函数
    ~SHADOW() {
        delete [] recv_pids;
        delete [] recv_offsets;
        if (subComm != MPI_COMM_NULL) MPI_Comm_free(& subComm);
        if (masterComm != MPI_COMM_NULL) MPI_Comm_free(& masterComm);
    }
};

/// @brief 聚合(Aggregation)类型的多重网格(MG, multigrid)
/// @tparam idx_t 索引类型
/// @tparam data_t 存储数据的浮点类型
/// @tparam calc_t 执行计算时的浮点类型
/// @tparam setup_t setup所用的浮点类型
template<typename idx_t, typename data_t, typename calc_t, typename setup_t>
class AggMG: public Precond<idx_t, calc_t, setup_t>
{
private:
    // setup时使用的线程数（实际未被用到）
    int setup_nt = 1;
    // 执行平滑计算的线程数（实际未被用到）
    int smth_nt = 1;

    // 执行粗化时生成的粗化层数
    idx_t coarse_nlevs = 0;
    // 记录生成粗化层后哪些原始粗化层会被跳过（即弃用）
    idx_t * skipped= nullptr;
    // 计算粗网格矩阵的三重矩阵乘积(RAP)时的系数
    calc_t* alphas = nullptr;
    // XPSS[i]记录了从第i层向第i+1层执行粗化时，进程数是否要减少
    // 如果值为1即不减少，否则以该值为步长执行进程数的缩减
    idx_t * XPSS = nullptr;
    // use_coord[i]记录了从从第i层向第i+1层执行粗化时，是否采用基于
    // 点坐标的粗化方法，值为1表示采用。否则采用基于最小堆的粗化算法
    idx_t * use_coord = nullptr;

    // 当前层是否可以执行零初值优化
    mutable bool zero_guess = false;
    // solve阶段的cycle类型：1 for V, 2 for W
    idx_t cycle_type = 1;
    // 外部传入的问题矩阵A，并不拥有，也不需释放
    const MyParCSRMatrix<idx_t, setup_t, setup_t> * base_A_ptr = nullptr;

    // 实际solve时的层数，刨除了那些被启用的粗化层
    idx_t num_levs = 0;
    // 多重网格是否拥有最细层的内存所有权
    bool own_A0 = false;
    // 每一层网格的矩阵
    MyParCSRMatrix<idx_t, data_t, calc_t> ** A_array = nullptr;
    // 每一层网格的向量
    MyParCSRVector<idx_t, calc_t> ** U_array = nullptr;
    MyParCSRVector<idx_t, calc_t> ** F_array = nullptr;
    MyParCSRVector<idx_t, calc_t> ** aux_arr = nullptr;
    // 每一层的smoother的类型
    std::string * relax_types = nullptr;
    // 每一层的smoother的权重系数
    calc_t * relax_wgts = nullptr;
    // 每一层的smoother
    Precond<idx_t, calc_t, setup_t> ** smoother = nullptr;

    // 每相邻两层间的插值算子
    idx_t ** plain_prlg = nullptr;
    // 每相邻两层间的限制算子
    MyCSRMatrix<idx_t, data_t, calc_t> ** restrictor = nullptr;

    // 在V-cycle中前平滑和后平滑的次数
    idx_t num_grid_sweeps[2] = {1, 1};

    // 用于记录在不同层间发生进程数收缩时的元数据
    SHADOW** shd_dict = nullptr;
    SHADOW** shd_setup= nullptr;

    /// @brief 执行第ilev层的前平滑
    /// @param ilev 层号
    void preSmooth(const idx_t ilev) const;
    /// @brief 执行第ilev层的后平滑
    /// @param ilev 层号
    void postSmooth(const idx_t ilev) const;
    /// @brief 在第ilev层内计算残差
    /// @param ilev 层号
    void calcResi(const idx_t ilev) const;
    /// @brief 将第ilev层的残差向第ilev+1层进行限制，并作为下一层求解的右端项
    /// @param ilev 层号
    void rstrResi(const idx_t ilev) const;
    /// @brief 将第ilev+1层的向量值插值回到第ilev层，以此更新第ilev层的向量
    /// @param ilev 层号
    void prlgErr(const idx_t ilev) const;
    /// @brief Solve阶段执行V-cycle过程
    void V_Cycle() const ;
    /// @brief Solve阶段执行W-cycle过程
    void W_Cycle() const ;

    /// @brief 处理配置参数的文件
    /// @param _config_file 文件名
    void ParseConfig(std::string _config_file);

    /// @brief 建立第ilev层的smoother
    /// @param ilev 层号
    /// @param A 第ilev层的矩阵
    void SetupSmoother(const idx_t ilev, MyParCSRMatrix<idx_t, setup_t, setup_t> & A);
    /// @brief 以stride为步长进行进程数的收缩。原来的0, stride-1, 2*stride-1, ..., 等进程会成为master进程，
    /// 而原来的其它进程则成为最近的（向上看齐）一个master进程的slave进程。一个master进程和stride-1个slave进程
    /// 组成一个收缩集合，其中的slave进程将目前负责的矩阵A的部分全部转移给master进程，由master进程继续进行更粗层
    /// 的计算（包括setup和solve）
    /// @param stride 收缩的步长，要求必须整除当前层的矩阵的通信域的进程数
    /// @param shd 记录收缩信息
    /// @param old_A 原来（收缩前）的矩阵
    /// @return 收缩后的矩阵，master进程的返回值非空，而slave进程返回为空
    MyParCSRMatrix<idx_t, setup_t, setup_t>* reducePx(const idx_t stride, SHADOW * & shd,
        const MyParCSRMatrix<idx_t, setup_t, setup_t> * old_A);
    /*
    void gather (const idx_t ilev) const;
    void scatter(const idx_t ilev) const;
    */
public:
    /// @brief 构造函数
    /// @param _config_file 参数配置文件的名字
    AggMG(const std::string _config_file);
    /// @brief 析构函数
    ~AggMG();
    /// @brief 建立函数
    /// @param A 调用者传进来的矩阵
    void Setup(MyParCSRMatrix<idx_t, setup_t, setup_t> & A);
    /// @brief 求解函数
    /// @param b 右端向量
    /// @param x 解向量
    /// @param _zg 是否采用零初值优化
    void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x, bool _zg=true);
};
typedef AggMG<int, double, double, double> AggMG_I32All64;
typedef AggMG<int, float , float , float > AggMG_I32All32;

#endif
